{
    "cells": [
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "import logging\n",
                "import sys\n",
                "\n",
                "logging.basicConfig(stream=sys.stdout, level=logging.INFO)\n",
                "logging.getLogger().addHandler(logging.StreamHandler(stream=sys.stdout))"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "from __future__ import absolute_import\n",
                "\n",
                "# My OpenAI Key\n",
                "import os\n",
                "os.environ['OPENAI_API_KEY'] = \"\"\n",
                "\n",
                "from llama_index.readers.database import DatabaseReader\n",
                "from llama_index import GPTSimpleVectorIndex"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "# Initialize DatabaseReader object with the following parameters:\n",
                "\n",
                "db = DatabaseReader(\n",
                "    scheme = \"postgresql\", # Database Scheme\n",
                "    host = \"localhost\", # Database Host\n",
                "    port = \"5432\", # Database Port\n",
                "    user = \"postgres\", # Database User\n",
                "    password = \"FakeExamplePassword\", # Database Password\n",
                "    dbname = \"postgres\", # Database Name\n",
                ")"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "### DatabaseReader class ###\n",
                "# db is an instance of DatabaseReader:\n",
                "print(type(db))\n",
                "# DatabaseReader available method:\n",
                "print(type(db.load_data))\n",
                "\n",
                "### SQLDatabase class ###\n",
                "# db.sql is an instance of SQLDatabase:\n",
                "print(type(db.sql_database))\n",
                "# SQLDatabase available methods:\n",
                "print(type(db.sql_database.from_uri))\n",
                "print(type(db.sql_database.get_single_table_info))\n",
                "print(type(db.sql_database.get_table_columns))\n",
                "print(type(db.sql_database.get_table_info))\n",
                "print(type(db.sql_database.get_table_names))\n",
                "print(type(db.sql_database.insert_into_table))\n",
                "print(type(db.sql_database.run))\n",
                "print(type(db.sql_database.run_sql))\n",
                "# SQLDatabase available properties:\n",
                "print(type(db.sql_database.dialect))\n",
                "print(type(db.sql_database.engine))\n",
                "print(type(db.sql_database.table_info))"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "### Testing DatabaseReader\n",
                "### from SQLDatabase, SQLAlchemy engine and Database URI:\n",
                "\n",
                "# From SQLDatabase instance:\n",
                "print(type(db.sql_database))\n",
                "db_from_sql_database = DatabaseReader(sql_database = db.sql_database)\n",
                "print(type(db_from_sql_database))\n",
                "\n",
                "# From SQLAlchemy engine:\n",
                "print(type(db.sql_database.engine))\n",
                "db_from_engine = DatabaseReader(engine = db.sql_database.engine)\n",
                "print(type(db_from_engine))\n",
                "\n",
                "# From Database URI:\n",
                "print(type(db.uri))\n",
                "db_from_uri = DatabaseReader(uri = db.uri)\n",
                "print(type(db_from_uri))"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "# The below SQL Query example returns a list values of each row\n",
                "# with concatenated text from the name and age columns\n",
                "# from the users table where the age is greater than or equal to 18\n",
                "\n",
                "query = f\"\"\"\n",
                "    SELECT\n",
                "        CONCAT(name, ' is ', age, ' years old.') AS text\n",
                "    FROM public.users\n",
                "    WHERE age >= 18\n",
                "    \"\"\""
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "# Please refer to llama_index.langchain_helpers.sql_wrapper\n",
                "# SQLDatabase.run_sql method\n",
                "texts = db.sql_database.run_sql(command = query)\n",
                "\n",
                "# Display type(texts) and texts\n",
                "# type(texts) must return <class 'list'>\n",
                "print(type(texts))\n",
                "\n",
                "# Documents must return a list of Tuple objects\n",
                "print(texts)"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "# Please refer to llama_index.readers.database.DatabaseReader.load_data\n",
                "# DatabaseReader.load_data method\n",
                "documents = db.load_data(query = query)\n",
                "\n",
                "# Display type(documents) and documents\n",
                "# type(documents) must return <class 'list'>\n",
                "print(type(documents))\n",
                "\n",
                "# Documents must return a list of Document objects\n",
                "print(documents)"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "try:\n",
                "    # Try to load existing Index from disk\n",
                "    index = GPTSimpleVectorIndex.load_from_disk('index.json')\n",
                "except:\n",
                "    index = GPTSimpleVectorIndex.from_documents(documents)\n",
                "\n",
                "    # Save newly created Index to disk\n",
                "    index.save_to_disk('index.json')"
            ]
        }
    ],
    "metadata": {
        "kernelspec": {
            "display_name": "Python 3 (ipykernel)",
            "language": "python",
            "name": "python3"
        },
        "language_info": {
            "codemirror_mode": {
                "name": "ipython",
                "version": 3
            },
            "file_extension": ".py",
            "mimetype": "text/x-python",
            "name": "python",
            "nbconvert_exporter": "python",
            "pygments_lexer": "ipython3",
            "version": "3.11.1"
        },
        "vscode": {
            "interpreter": {
                "hash": "bd5508c2ffc7f17f7d31cf4086cc872f89e96996a08987e995649e5fbe85a3a4"
            }
        }
    },
    "nbformat": 4,
    "nbformat_minor": 2
}